"""
Run a rest API exposing the yolov5s object detection model
"""
import argparse
import io
import json
import numpy as np

import base64
import cv2
from io import BytesIO
from datetime import date
import datetime
import os

import torch
from PIL import Image
from flask import Flask, request, send_from_directory
from cfg import *
import logging

app = Flask(__name__)
logging.basicConfig(filename='log.log', level=logging.DEBUG, format='%(asctime)s %(levelname)s %(name)s %(threadName)s : %(message)s')

DETECTION_URL = "/v1/object-detection/person-car"
IMAGE_RESULT = "/results"
# ResultSaveDays = 3
# # 电子围栏构建
# ELECTRONIC_FENCE = {
#     # key   name; Value     PointsList
#     'ljk1': {
#         'Description': "dam name is ljk1",
#         'PointsList': [[5, 446 ], [3, 500 ], [1413.941176470588, 918.1176470588234 ], [1919, 897 ], [1919, 394 ], [1124, 355 ], [1124, 355 ], [456, 396 ], [367, 399 ], [311, 400 ], [210, 409 ] ]
#     },
#     'sn1': {
#         'Description': "dam name is sn1",
#         'PointsList': [[959, 753 ], [958, 776 ], [1099, 824 ], [1070, 829 ], [1081, 841 ], [1068, 846 ], [1091, 854 ], [1177, 858 ], [1220, 873 ], [1359, 912 ], [1356, 925 ], [1244, 955 ], [1546, 1054 ], [1705, 1328 ], [1754, 1286 ], [1841, 1309 ], [1867, 1388 ], [2179, 1439 ], [2556, 1436 ], [2556, 771 ], [1951, 761 ], [1839, 739 ], [1721, 742 ], [1475, 762 ], [1076, 759 ] ]
#     }
# }

@app.route(DETECTION_URL, methods=["POST"])
def predict():
    if request.files.get("image"):
        image_file = request.files["image"]
        # 读取二进制不成功或者图片打开不成功两种可能性
        data_json = {}
        try:
            image_bytes = image_file.read()
            img = Image.open(io.BytesIO(image_bytes))
        except:
            data_json['erro_code'] = 1100
            data_json['message'] = '读取二进制不成功或者图片打开不成功'
            return data_json
        try:
            results = model(img, size=640)  # reduce size=320 for faster inference
        except:
            data_json['erro_code'] = 2200
            data_json['message'] = '模型预测不成功'
            return data_json
        try:
            # updates results.imgs with boxes and labels, returns nothing
            results.render()
        except:
            data_json['erro_code'] = 2201
            data_json['message'] = '预测结果渲染不成功'
            return data_json
        # info details summary
        try:
            data_json['summary']  = json.loads(results.pandas().xyxy[0]['name'].value_counts().to_json())
            data_json['details'] = json.loads(results.pandas().xyxy[0].to_json(orient="records"))
        except:
            data_json['erro_code'] = 2202
            data_json['message'] = 'Summary和Details结果为空'
            return data_json

        # 结果正常
        data_json['erro_code'] = 0
        data_json['message'] = 'OK'
        # 只要结果渲染成功就可以返回URL
        if request.values.get('Base64') == 'True':
            for img in results.imgs:
                buffered = BytesIO()
                img_base64 = Image.fromarray(img)
                img_base64.save(buffered, format="JPEG")
                data_json['Base64'] = base64.b64encode(buffered.getvalue()).decode('utf-8')
        else:
            data_json['Base64'] = None
        # 只要结果渲染成功就可以返回URL
        
        if request.values.get('RenderURL') == 'True' and request.values.get('ImageName'):
            # 2019-12-11
            today = date.today()
            img_path = IMAGE_RESULT + f"/{today}/"
            os.makedirs("/app" + img_path, exist_ok=True)
            for img in results.imgs:
                img_res = Image.fromarray(img)
                img_res.save("/app" + img_path + request.values['ImageName'])
                data_json['RenderURL'] = img_path + request.values['ImageName']
        else:
            data_json['RenderURL'] = None

        return data_json

@app.route(f'{IMAGE_RESULT}/<path:imagename>', methods=["GET", "POST"])
def get_image_result(imagename):
    """IMAGE_RESULT/date/imagename
    """
    # today = date.today()
    gen_date_str = os.path.dirname(imagename)
    imagename = os.path.basename(imagename)
    return send_from_directory("/app" + IMAGE_RESULT + "/" + gen_date_str, imagename, as_attachment=False)

@app.route('/v1/eletronic/getall', methods=['POST'])
def get_all_eletronic_fence():
    # with description
    return ELECTRONIC_FENCE

@app.route('/v1/electronic/get', methods=['POST'])
def get_the_electronic_fence():
    # Post的时候，要求，data=json.dumps(data)
    try:
        data_dict = json.loads(request.data)
    except:
        return {'erro_code': 2202, 'message': 'post requst data is None'}
    if not data_dict.get('ele_fence_key'):
        return {'erro_code': 3300, 'message': '参数key缺失'}
    ele_fence_key = data_dict.get('ele_fence_key')
    if not ELECTRONIC_FENCE.get(ele_fence_key):
        return {'erro_code': 3301, 'message': '参数key存在，但参数value缺失'}
    data_json = ELECTRONIC_FENCE.get(ele_fence_key)
    if isinstance(data_json, dict):
        return data_json
    else:
        return {'erro_code': 3302, 'message': '参数key存在，参数value也存在，但是数据格式不合规'}

@app.route('/v1/electronic/update-or-add', methods=['POST'])
def update_electronic_fence():
    # Post的时候，要求，data=json.dumps(data)
    try:
        data_dict = json.loads(request.data)
    except:
        return {'erro_code': 2202, 'message': 'post requst data is None'}
    if not data_dict.get('ele_fence_key'):
        return {'erro_code': 3300, 'message': '参数key缺失'}
    ele_fence_key = data_dict.get('ele_fence_key')
    if not data_dict.get('ele_fence_value'):
        return {'erro_code': 3300, 'message': '参数key缺失'}
    # 直接赋值就好，不管是否为空。
    ELECTRONIC_FENCE[ele_fence_key] = data_dict.get('ele_fence_value')
    # 赋值，不做数值的校验
    return {'erro_code': 0, 'message': 'OK'}

def init_electronic_fence(ele_fence_key):
    if not ELECTRONIC_FENCE.get(ele_fence_key):
        # 无电子围栏，请使用，无电子围栏接口
        return None, None
    else:
        ele_fence = ELECTRONIC_FENCE.get(ele_fence_key)
        # 根据视频尺寸，填充一个polygon，供撞线计算使用
        # lk1 1920 1080
        mask_image = np.zeros(ele_fence['OriginSize'], dtype=np.uint8)
        # 初始化lk1的polygon
        ndarray_pts = np.array(ele_fence['PointsList'], np.int32)
        polygon_blue = cv2.fillPoly(mask_image, [ndarray_pts], color=1)
        polygon_blue = polygon_blue[:, :, np.newaxis]
        #
        polygon_mask_blue = cv2.resize(polygon_blue, ele_fence['MaskResize'])
        #
        blue_color_plate = [255, 0, 0]
        # 蓝 polygon图片
        # 里面只有0和2555
        blue_color_image = np.array(polygon_blue * blue_color_plate, np.uint8)
        blue_color_image = cv2.resize(blue_color_image, ele_fence['MaskResize'])
        return polygon_mask_blue, blue_color_image, ele_fence['MaskResize']

@app.route(ELECTRONIC_FENCE_ABNORMAL_URL, methods=["POST"])
def inside_abnormal():
    if request.files.get("image"):
        image_file = request.files["image"]
        # 读取二进制不成功或者图片打开不成功两种可能性
        data_json = {}
        try:
            image_bytes = image_file.read()
            im_origin = np.array(Image.open(io.BytesIO(image_bytes)).convert('RGB'))
            im_origin = im_origin[:, :, ::-1].copy()
        except:
            data_json['erro_code'] = 1100
            data_json['message'] = '读取二进制不成功或者图片打开不成功'
            return data_json
        if request.values.get('ele_fence_key'):
            # 进行初始化
            polygon_mask_blue, blue_color_image, mask_resize = init_electronic_fence(request.values['ele_fence_key'])
            if polygon_mask_blue is None or blue_color_image is None:
                data_json['erro_code'] = 3309
                data_json['message'] = '无电子围栏设置，请设置电子围栏'
                return data_json
        else:
            data_json['erro_code'] = 3300
            data_json['message'] = '参数key缺失'
            return data_json
        try:
            # 坐标对齐这句太关键了。
            im = cv2.resize(im_origin, mask_resize)
            # 这里用的模型是modelx6精度最高的模型。
            results = modelx6(im.copy(), size=640)  # reduce size=320 for faster inference
        except:
            data_json['erro_code'] = 2200
            data_json['message'] = '模型预测不成功'
            return data_json
        # reach here electronic fence setup over
        # set up the render
        font_draw_number = cv2.FONT_HERSHEY_SIMPLEX
        draw_text_postion = (399, 48)
        output_image_frame = cv2.add(im, blue_color_image)
        # 1、获取到识别结果
        # 2、计算是否相交
        # 3、对相交区域进行色块填充
        # 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name'  # xyxy columns
        list_bboxs = results.pandas().xyxy[0].values.tolist()
        if len(list_bboxs) > 0:
            #
            inside_count = {}
            # inside_details
            inside_details = []
            for item_bbox in list_bboxs:
                x1, y1, x2, y2, confidence, class_id, label = item_bbox
                x1 = int(x1); y1 = int(y1); x2 = int(x2); y2 = int(y2)
                y1_offset = int(y1 + ((y2 - y1) * 0.618))
                # 撞线的点
                y = y1_offset
                x = x1
                if (polygon_mask_blue[y, x] == 1).any():
                    # 加入details列表
                    inside_details.append(item_bbox)
                    # 如果撞蓝polygon
                    inside_count['Abnormal'] = inside_count.get('Abnormal', 0) + 1
                    output_image_frame = cv2.drawContours(output_image_frame, [np.array([x1, y1, x1, y2, x2, y2, x2, y1]).reshape((-1, 2))], -1, (0, 0, 255), 4)
            # 所有检测内容碰撞完毕
            # 进行最终结果渲染。
            # text_list
            text_list = []
            text_list.append('INSIDE: Electronic Fence is working')
            if len(inside_count) >= 1:
                text_list.extend(["    " + str(key) + ":" + str(value) for key, value in inside_count.items()])
                # 通过焦点数量判断是否包含异常
                data_json['AbnormalDetails'] = inside_details
                data_json['IsAbnormal'] = True
                data_json['AbnormalSummary'] = inside_count['Abnormal']
            else:
                data_json['AbnormalDetails'] = None
                data_json['IsAbnormal'] = False
                data_json['AbnormalSummary'] = None
                text_list.append("    Everything Is Normal")
            # 把所有文本渲染到图片上
            y0, dy = draw_text_postion[1], 36
            for i, line in enumerate(text_list):
                y = y0 + i*dy
                output_image_frame = cv2.putText(img=output_image_frame, text=line, org=(draw_text_postion[0], y), fontFace=font_draw_number, fontScale=1, color=(0, 0, 255), thickness=2)
        else:
            # 模型未识别到任何异常或者非异常的物体
            data_json['erro_code'] = 2209
            data_json['message'] = '模型未识别到任何异常或者非异常的物体'
            return data_json
            # data_json['AbnormalInside'] = False

        # 结果正常
        data_json['erro_code'] = 0
        data_json['message'] = 'OK'
        # 只要结果渲染成功就可以返回URL
        if request.values.get('Base64') == 'True':
            buffered = BytesIO()

            output_image_frame = cv2.cvtColor(output_image_frame, cv2.COLOR_BGR2RGB)

            img_base64 = Image.fromarray(output_image_frame)
            img_base64.save(buffered, format="JPEG")
            data_json['Base64'] = base64.b64encode(buffered.getvalue()).decode('utf-8')
        else:
            data_json['Base64'] = None
        # 只要结果渲染成功就可以返回URL

        if request.values.get('RenderURL') == 'True' and request.values.get('ImageName'):
            # 2019-12-11
            today = date.today()
            img_path = IMAGE_RESULT + f"/ele_fence/{today}/"
            os.makedirs("/app" + img_path, exist_ok=True)

            output_image_frame = cv2.cvtColor(output_image_frame, cv2.COLOR_BGR2RGB)

            img_res = Image.fromarray(output_image_frame)
            img_res.save("/app" + img_path + request.values['ImageName'])
            data_json['RenderURL'] = img_path + request.values['ImageName']
        else:
            data_json['RenderURL'] = None

        return data_json




if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Person and Car Usage")
    parser.add_argument("--port", default=8501, type=int, help="port number")
    parser.add_argument("--model_type", default='m6', help="model name")
    parser.add_argument("--conf", default=0.1, type=float, help="confidence threshold")
    parser.add_argument("--iou", default=0.65, type=float, help="iou threshold")
    parser.add_argument("--classes", default=[0, 1, 2, 3, 5, 7, 8], type=float, help="person bicycle car motorcycle bus truck boat")
    args = parser.parse_args()

    # model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True)  # force_reload to recache
    model = torch.hub.load('/app/weights_or_data/', 'yolov5{}'.format(args.model_type), pretrained=True, source='local')
    model.conf = args.conf
    model.iou = args.iou
    model.classes = args.classes
    # modelx6系列解决电子围栏里面的可疑物体。
    modelx6 = torch.hub.load('/app/weights_or_data/', 'yolov5x6', pretrained=True, source='local')
    modelx6.conf = 0.01
    modelx6.iou = 0.65

    app.run(host="0.0.0.0", port=args.port, debug=False)  # debug=True causes Restarting with stat
